import os
import cv2
import numpy as np
from tqdm import tqdm  # 引入进度条库

# 输入和输出路径 context59
maskclip_dir = "/space0/liangc/code/study/OVSS/MaskCLIP/output/voc59"
ours_dir = "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/demo/voc59_dclip"
gt_dir = "/space0/liangc/code/study/OVSS/MaskCLIP/output/voc59gt"
img_dir = "/space0/liangc/code/study/OVSS/MaskCLIP/data/VOCdevkit/VOC2010/JPEGImages"
output_dir = "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/demo/voc59_compare4"

# cityscapes
# maskclip_dir = "/space0/liangc/code/study/OVSS/MaskCLIP/output/city_maskclip/frankfurt"
# ours_dir = "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/demo/cityscapes_dclip/frankfurt"
# gt_dir = "/space0/liangc/code/study/OVSS/MaskCLIP/gt/cityscapes/frankfurt"
# img_dir = "/space0/liangc/code/study/OVSS/MaskCLIP/images/cityscapes/frankfurt"
# output_dir = "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/demo/cityscapes_compare/frankfurt"

os.makedirs(output_dir, exist_ok=True)

# 获取文件列表，假设四个文件夹内文件名完全一致
file_list = sorted(os.listdir(maskclip_dir))

# 使用 tqdm 进度条
for file_name in tqdm(file_list, desc="生成对比图", unit="张"):
    maskclip_path = os.path.join(maskclip_dir, file_name)
    ours_path = os.path.join(ours_dir, file_name)
    gt_path = os.path.join(gt_dir, file_name)
    img_path = os.path.join(img_dir, file_name)

    # 检查四个文件是否都存在
    if not all(os.path.exists(p) for p in [maskclip_path, ours_path, gt_path, img_path]):
        print(f"文件缺失：{file_name}")
        continue

    # 读取图片
    maskclip_img = cv2.imread(maskclip_path)
    ours_img = cv2.imread(ours_path)
    gt_img = cv2.imread(gt_path)
    img = cv2.imread(img_path)

    # 调整为相同大小（取 MaskCLIP 的大小为基准）
    height, width = maskclip_img.shape[:2]
    ours_img = cv2.resize(ours_img, (width, height))
    gt_img = cv2.resize(gt_img, (width, height))
    img = cv2.resize(img, (width, height))

    # 水平拼接两两子图
    upper_row = np.hstack((img, gt_img))
    lower_row = np.hstack((maskclip_img, ours_img))
    combined_img = np.vstack((upper_row, lower_row))

    # 保存拼接图片
    output_path = os.path.join(output_dir, file_name)
    cv2.imwrite(output_path, combined_img)

print("图片对比图生成完毕！")
